import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from statsmodels.tsa.seasonal import seasonal_decompose

class SeasonalityAnalyzer:
    def __init__(self, data=None):
        self.data = data
        self.decomposition = None
        
    def analyze_seasonality(self, data=None, column='sales_quantity', period=7):
        """分析时间序列的季节性"""
        if data is not None:
            self.data = data
            
        # 确保数据是时间序列格式
        if not isinstance(self.data.index, pd.DatetimeIndex):
            raise ValueError("数据索引必须是日期时间类型")
            
        # 检查数据是否足够长
        if len(self.data) < 2 * period:
            raise ValueError(f"数据长度({len(self.data)})不足以进行季节性分解(至少需要{2 * period}个观测值)")
            
        # 季节性分解
        self.decomposition = seasonal_decompose(
            self.data[column], 
            model='multiplicative', 
            period=period
        )
        
        return self.decomposition
    
    def plot_seasonality(self, save_path=None):
        """绘制季节性分解结果"""
        if self.decomposition is None:
            raise ValueError("必须先进行季节性分解")
            
        fig, (ax1, ax2, ax3, ax4) = plt.subplots(4, 1, figsize=(12, 10))
        
        # 原始数据
        self.decomposition.observed.plot(ax=ax1)
        ax1.set_title('原始销售数据')
        ax1.set_xlabel('')
        
        # 趋势
        self.decomposition.trend.plot(ax=ax2)
        ax2.set_title('销售趋势')
        ax2.set_xlabel('')
        
        # 季节性
        self.decomposition.seasonal.plot(ax=ax3)
        ax3.set_title('季节性模式')
        ax3.set_xlabel('')
        
        # 残差
        self.decomposition.resid.plot(ax=ax4)
        ax4.set_title('残差')
        
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path)
            
        return fig
    
    def get_day_of_week_pattern(self):
        """获取每周内各天的销售模式"""
        if self.data is None:
            raise ValueError("必须先加载数据")
            
        # 添加星期几信息
        data_copy = self.data.copy()
        if isinstance(data_copy, pd.DataFrame) and not isinstance(data_copy.index, pd.DatetimeIndex):
            if 'date' in data_copy.columns:
                data_copy['weekday'] = data_copy['date'].dt.dayofweek
        else:
            data_copy['weekday'] = data_copy.index.dayofweek
            
        # 计算每周各天的平均销售量
        weekday_avg = data_copy.groupby('weekday')['sales_quantity'].mean().reindex(range(7))
        
        # 映射到星期名称
        weekday_names = ['周一', '周二', '周三', '周四', '周五', '周六', '周日']
        weekday_avg.index = weekday_names
        
        return weekday_avg
    
    def get_monthly_pattern(self):
        """获取每月的销售模式"""
        if self.data is None:
            raise ValueError("必须先加载数据")
            
        # 添加月份信息
        data_copy = self.data.copy()
        if isinstance(data_copy, pd.DataFrame) and not isinstance(data_copy.index, pd.DatetimeIndex):
            if 'date' in data_copy.columns:
                data_copy['month'] = data_copy['date'].dt.month
                # 添加年份信息，用于检查数据跨度
                data_copy['year'] = data_copy['date'].dt.year
        else:
            data_copy['month'] = data_copy.index.month
            # 添加年份信息，用于检查数据跨度
            data_copy['year'] = data_copy.index.year
        
        # 检查数据是否跨越至少12个月
        unique_year_months = data_copy[['year', 'month']].drop_duplicates()
        months_count = len(unique_year_months)
        
        if months_count < 12:
            print(f"警告：数据仅包含{months_count}个月，月度季节性分析需要至少12个月的数据")
            
        # 检查是否有足够的年份
        years_count = data_copy['year'].nunique()
        if years_count < 2:
            print(f"警告：数据仅跨越{years_count}年，月度季节性分析更准确需要至少2年数据")
            
        # 计算每月的平均销售量
        monthly_avg = data_copy.groupby('month')['sales_quantity'].mean().reindex(range(1, 13))
        
        # 映射到月份名称
        month_names = ['1月', '2月', '3月', '4月', '5月', '6月', '7月', '8月', '9月', '10月', '11月', '12月']
        monthly_avg.index = month_names
        
        return monthly_avg
